import torch
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
import torch.utils.data as torchdata
# from Retrieval_tools import *

def default_loader(path):
    return Image.open(path).convert('RGB')


class MyDataset(torchdata.Dataset):
    def __init__(self, txt='../dataset/VOCdevkit/onehot_trainval_2007',
                 transform=None, target_transform=None, loader=default_loader):
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader
        file = torch.load(txt)
        imgs = []
        for i in range(len(file)):
        # for i in range(10):
            line = file[i]
            imgpath = line.strip().split()[0]
            label = line.strip().split()[1:]
            label = list(map(int, label))
            words = [imgpath, label]
            imgs.append(words)
        self.imgs = imgs

    def __getitem__(self, index):
        words = self.imgs[index]
        img = self.loader('../dataset/VOCdevkit/VOC2007/JPEGImages/'+words[0]+'.jpg')
        if self.transform is not None:
            img = self.transform(img)
        label = torch.LongTensor(words[1])
        return img, label
        # ----------- 以下为修改部分 -------------------------------------------------------
        # tem = []
        # tem.append(int(words[0]))
        # img_name = torch.Tensor(tem).int()
        # # print(img_name)
        # return img_name, img, label
        # ---------- end ------------------------------------------------------------------

    def __len__(self):
        return len(self.imgs)

def init_voc_dataloader(args):

    # mean_value = (104, 117, 123)
    # mean_value = (0.4078, 0.4588, 0.4824)

    mean = (0.4914, 0.4822, 0.4465)
    std = (0.2023, 0.1994, 0.2010)
    transform_train = transforms.Compose([
        # transforms.Resize((224, 224)),
        transforms.Resize((448, 448)),
        # transforms.RandomCrop((448, 448)),
        # transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])
    transform_test = transforms.Compose([
        transforms.Resize((448, 448)),
        # transforms.RandomCrop((448, 448)),
        # transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])

    test_dir = '../dataset/VOCdevkit/onehot_test_2007'
    trainval_dir = '../dataset/VOCdevkit/onehot_trainval_2007'

    test_set = MyDataset(txt=test_dir, transform=transform_test)
    test_loader = torchdata.DataLoader(test_set, batch_size=args.batch_size, shuffle=True, num_workers=4)
   
    trainval_set = MyDataset(txt=trainval_dir, transform=transform_train)
    trainval_loader = torchdata.DataLoader(trainval_set, batch_size=args.batch_size, shuffle=True, num_workers=4)
    
    compute_test_loader = torchdata.DataLoader(test_set, batch_size=10, shuffle=True, num_workers=4)
    compute_trainval_loader = torchdata.DataLoader(trainval_set, batch_size=10, shuffle=True, num_workers=4)
    
    return trainval_loader, test_loader, compute_test_loader, compute_trainval_loader

def init_voc_dataloader_mirror(args):

    # mean_value = (104, 117, 123)
    # mean_value = (0.4078, 0.4588, 0.4824)

    mean = (0.4914, 0.4822, 0.4465)
    std = (0.2023, 0.1994, 0.2010)
    transform_train = transforms.Compose([
        # transforms.Resize((224, 224)),
        transforms.Resize((448, 448)),
        # transforms.RandomCrop((448, 448)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])
    transform_test = transforms.Compose([
        transforms.Resize((448, 448)),
        # transforms.RandomCrop((448, 448)),
        # transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])

    test_dir = '../dataset/VOCdevkit/onehot_test_2007'
    trainval_dir = '../dataset/VOCdevkit/onehot_trainval_2007'

    trainval_set = MyDataset(txt=trainval_dir, transform=transform_train)
    trainval_loader = torchdata.DataLoader(trainval_set, batch_size=args.batch_size, shuffle=True, num_workers=4)


    test_set = MyDataset(txt=test_dir, transform=transform_test)
    test_loader = torchdata.DataLoader(test_set, batch_size=10, shuffle=False, num_workers=4)
    db_set = MyDataset(txt=trainval_dir, transform=transform_test)
    db_loader = torchdata.DataLoader(db_set, batch_size=10, shuffle=False, num_workers=4)
    
    
    # compute_test_loader = torchdata.DataLoader(test_set, batch_size=10, shuffle=True, num_workers=4)
    # compute_trainval_loader = torchdata.DataLoader(trainval_set, batch_size=10, shuffle=True, num_workers=4)
    
    return trainval_loader, test_loader, db_loader # , compute_test_loader, compute_trainval_loader

def init_voc_dataloader_mirror_2(args):

    # mean_value = (104, 117, 123)
    # mean_value = (0.4078, 0.4588, 0.4824)

    mean = (0.4914, 0.4822, 0.4465)
    std = (0.2023, 0.1994, 0.2010)
    transform_train = transforms.Compose([
        # transforms.Resize((224, 224)),
        transforms.Resize((448, 448)),
        # transforms.RandomCrop((448, 448)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])
    transform_test = transforms.Compose([
        transforms.Resize((448, 448)),
        # transforms.RandomCrop((448, 448)),
        # transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])

    test_dir = '../dataset/VOCdevkit/onehot_test_2007'
    trainval_dir = '../dataset/VOCdevkit/onehot_trainval_2007'

    trainval_set = MyDataset(txt=trainval_dir, transform=transform_train)
    trainval_loader = torchdata.DataLoader(trainval_set, batch_size=args.batch_size, shuffle=True, num_workers=4)


    test_set = MyDataset(txt=test_dir, transform=transform_test)
    test_loader = torchdata.DataLoader(test_set, batch_size=10, shuffle=False, num_workers=4)
   
    
    
    # compute_test_loader = torchdata.DataLoader(test_set, batch_size=10, shuffle=True, num_workers=4)
    # compute_trainval_loader = torchdata.DataLoader(trainval_set, batch_size=10, shuffle=True, num_workers=4)
    
    return trainval_loader, test_loader # , compute_test_loader, compute_trainval_loader



def init_voc_dataloader_extend(args):

    # mean_value = (104, 117, 123)
    # mean_value = (0.4078, 0.4588, 0.4824)

    mean = (0.4914, 0.4822, 0.4465)
    std = (0.2023, 0.1994, 0.2010)
    transform_train = transforms.Compose([
        # transforms.Resize((224, 224)),
        transforms.Resize((512, 512)),
        transforms.RandomCrop((448, 448)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])
    transform_test = transforms.Compose([
        transforms.Resize((512, 512)),
        transforms.RandomCrop((448, 448)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])

    test_dir = '../dataset/VOCdevkit/onehot_test_2007'
    trainval_dir = '../dataset/VOCdevkit/onehot_trainval_2007'

    test_set = MyDataset(txt=test_dir, transform=transform_test)
    test_loader = torchdata.DataLoader(test_set, batch_size=args.batch_size, shuffle=True, num_workers=4)
   
    trainval_set = MyDataset(txt=trainval_dir, transform=transform_train)
    trainval_loader = torchdata.DataLoader(trainval_set, batch_size=args.batch_size, shuffle=True, num_workers=4)
    
    compute_test_loader = torchdata.DataLoader(test_set, batch_size=10, shuffle=True, num_workers=4)
    compute_trainval_loader = torchdata.DataLoader(trainval_set, batch_size=10, shuffle=True, num_workers=4)
    
    return trainval_loader, test_loader, compute_test_loader, compute_trainval_loader
    

categories = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair',
         'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant',
         'sheep', 'sofa', 'train', 'tvmonitor']





# --------------------------------------------------------------------------------------


class MyDataset_testing(torchdata.Dataset):
    def __init__(self, txt='../dataset/VOCdevkit/onehot_trainval_2007',
                 transform=None, target_transform=None, loader=default_loader):
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader
        file = torch.load(txt)
        imgs = []
        # for i in range(len(file)):
        for i in range(48):
            line = file[i]
            imgpath = line.strip().split()[0]
            label = line.strip().split()[1:]
            label = list(map(int, label))
            words = [imgpath, label]
            imgs.append(words)
        self.imgs = imgs

    def __getitem__(self, index):
        words = self.imgs[index]
        img = self.loader('../dataset/VOCdevkit/VOC2007/JPEGImages/'+words[0]+'.jpg')
        if self.transform is not None:
            img = self.transform(img)
        label = torch.LongTensor(words[1])
        return img, label
        # ----------- 以下为修改部分 -------------------------------------------------------
        # tem = []
        # tem.append(int(words[0]))
        # img_name = torch.Tensor(tem).int()
        # # print(img_name)
        # return img_name, img, label
        # ---------- end ------------------------------------------------------------------

    def __len__(self):
        return len(self.imgs)

# transforms.Resize((256, 256)),
# transforms.RandomCrop((224, 224)),
# transforms.RandomHorizontalFlip(),
# transforms.ToTensor(),
# normalize


def init_voc_dataloader_testing(args):

    # mean_value = (104, 117, 123)
    # mean_value = (0.4078, 0.4588, 0.4824)

    mean = (0.4914, 0.4822, 0.4465)
    std = (0.2023, 0.1994, 0.2010)
    # mean = (0.4914, 0.4822, 0.4465)
    # std = (0.2023, 0.1994, 0.2010)
    transform_train = transforms.Compose([
        # transforms.Resize((224, 224)),
        transforms.Resize((448, 448)),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])
    transform_test = transforms.Compose([
        # transforms.Resize((224, 224)),
        transforms.Resize((448, 448)),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])

    test_dir = '../dataset/VOCdevkit/onehot_test_2007'
    trainval_dir = '../dataset/VOCdevkit/onehot_trainval_2007'

    test_set = MyDataset_testing(txt=test_dir, transform=transform_test)
    test_loader = torchdata.DataLoader(test_set, batch_size=args.batch_size, shuffle=True, num_workers=4)
   
    trainval_set = MyDataset_testing(txt=trainval_dir, transform=transform_train)
    trainval_loader = torchdata.DataLoader(trainval_set, batch_size=args.batch_size, shuffle=True, num_workers=4)
    
    compute_test_loader = torchdata.DataLoader(test_set, batch_size=10, shuffle=True, num_workers=4)
    compute_trainval_loader = torchdata.DataLoader(trainval_set, batch_size=10, shuffle=True, num_workers=4)
    
    return trainval_loader, test_loader, compute_test_loader, compute_trainval_loader
    # return DB_loader, trainval_loader, test_loader


